"""
信号发生器模块
基于RsSmw_FMCW.py重构，负责SMW信号发生器的控制和FMCW信号生成
"""

import numpy as np
import tempfile
import os
from typing import Tuple, Dict, Any, Optional
from RsSmw import *
from src.logger import Logger


class SignalGenerator:
    """SMW信号发生器控制器"""
    
    def __init__(self, config: Dict[str, Any], logger: Optional[Logger] = None):
        """
        初始化信号发生器
        
        Args:
            config: SMW仪器配置
            logger: 日志记录器
        """
        self.config = config
        self.logger = logger or Logger()
        self.smw = None
        self.pc_wv_file = None
        self.instr_wv_file_outA = '/var/user/InstrDemoFile_outA.wv'
        self.instr_wv_file_outB = '/var/user/InstrDemoFile_outB.wv'
        
    def connect(self) -> bool:
        """
        连接SMW信号发生器
        
        Returns:
            连接是否成功
        """
        try:
            ip = self.config.get('ip', '192.168.31.235')
            self.logger.info(f"正在连接SMW信号发生器: {ip}")
            
            # 检查RsSmw模块版本
            RsSmw.assert_minimum_version('5.0.44')
            
            # 建立连接
            self.smw = RsSmw(f'TCPIP::{ip}::hislip0')
            
            # 获取仪器信息
            idn_string = self.smw.utilities.idn_string
            self.logger.info(f"SMW仪器信息: {idn_string}")
            
            # 重置仪器到已知状态
            #self.smw.utilities.reset()
            self.logger.info("SMW仪器已重置")
            
            return True
            
        except Exception as e:
            self.logger.error(f"连接SMW信号发生器失败: {e}")
            return False
    
    def generate_fmcw_iq_data(self, fmcw_config: Dict[str, Any]) -> Tuple[np.ndarray, np.ndarray]:
        """
        生成FMCW信号的I/Q数据
        
        Args:
            fmcw_config: FMCW信号配置
            
        Returns:
            (I数据, Q数据) 元组
        """
        # 提取配置参数
        center_freq = self.config.get('center_frequency', 10e9)
        clock_freq = self.config.get('clock_frequency', 600e6)
        
        bandwidth = fmcw_config.get('bandwidth', 100e6)
        sweep_rate = fmcw_config.get('sweep_rate', 2e12)
        scale_factor = fmcw_config.get('scale_factor', 0.8)
        
        # 新增参数：三角波配置
        waveform_type = fmcw_config.get('waveform_type', 'linear')
        down_bandwidth = fmcw_config.get('down_bandwidth', 100e6)
        down_sweep_rate = fmcw_config.get('down_sweep_rate', 20e12)
        
        self.logger.info("生成FMCW信号参数:")
        self.logger.info(f"  中心频率: {center_freq/1e9:.3f} GHz")
        self.logger.info(f"  波形类型: {waveform_type}")
        self.logger.info(f"  时钟频率: {clock_freq/1e6:.1f} MHz")
        self.logger.info(f"  缩放因子: {scale_factor}")
        
        if waveform_type == 'triangular':
            # 三角波配置
            self.logger.info(f"  下降部分带宽: {down_bandwidth/1e6:.1f} MHz")
            self.logger.info(f"  下降部分扫频速率: {down_sweep_rate/1e12:.1f} MHz/μs")
            
            # 计算三角波的时间参数
            down_duration = down_bandwidth / down_sweep_rate
            up_duration = down_bandwidth / sweep_rate  # 假设上升部分使用相同的带宽
            total_duration = down_duration + up_duration
            
            # 生成时间向量
            time_vector = np.arange(0, total_duration, 1 / clock_freq)
            
            # 生成三角波的瞬时频率
            instantaneous_freq = np.zeros_like(time_vector)
            
            # 下降部分：从 +bandwidth/2 到 -bandwidth/2
            down_samples = int(down_duration * clock_freq)
            down_indices = time_vector < down_duration
            instantaneous_freq[down_indices] = down_bandwidth/2 - down_sweep_rate * time_vector[down_indices]
            
            # 上升部分：从 -bandwidth/2 到 +bandwidth/2
            up_indices = time_vector >= down_duration
            up_time = time_vector[up_indices] - down_duration
            instantaneous_freq[up_indices] = -down_bandwidth/2 + sweep_rate * up_time
            
            self.logger.info(f"  下降时间: {down_duration*1e6:.1f} μs")
            self.logger.info(f"  上升时间: {up_duration*1e6:.1f} μs")
            self.logger.info(f"  总时间: {total_duration*1e6:.1f} μs")
            self.logger.info(f"  频率范围: {-down_bandwidth/2/1e6:.1f} MHz 到 {down_bandwidth/2/1e6:.1f} MHz")
            
        else:
            # 原始线性扫频
            self.logger.info(f"  扫频速率: {sweep_rate/1e12:.1f} MHz/μs")
            self.logger.info(f"  带宽: {bandwidth/1e6:.1f} MHz")
            
            # 计算信号参数
            sweep_duration = bandwidth / sweep_rate
            time_vector = np.arange(0, sweep_duration, 1 / clock_freq)
            
            # 生成线性扫频
            instantaneous_freq = -bandwidth/2 + sweep_rate * time_vector
            
            self.logger.info(f"  基带频率扫描范围: {-bandwidth/2/1e6:.1f} MHz 到 {bandwidth/2/1e6:.1f} MHz")
        
        # 相位是频率的积分 (基带)
        phase = 2 * np.pi * np.cumsum(instantaneous_freq) / clock_freq
        
        # I分量和Q分量数据
        i_data = np.cos(phase) * scale_factor
        q_data = np.sin(phase) * scale_factor
        
        # 记录统计信息
        self.logger.info("I/Q数据统计:")
        self.logger.info(f"  I数据 - 范围: [{np.min(i_data):.4f}, {np.max(i_data):.4f}], "
                        f"均值: {np.mean(i_data):.4f}, 标准差: {np.std(i_data):.4f}")
        self.logger.info(f"  Q数据 - 范围: [{np.min(q_data):.4f}, {np.max(q_data):.4f}], "
                        f"均值: {np.mean(q_data):.4f}, 标准差: {np.std(q_data):.4f}")
        self.logger.info(f"  采样点数: {len(time_vector)}")
        
        if waveform_type == 'triangular':
            actual_bandwidth = down_bandwidth
        else:
            actual_bandwidth = bandwidth
            
        self.logger.info(f"  实际RF频率扫描范围: {(center_freq-actual_bandwidth/2)/1e9:.3f} GHz 到 "
                        f"{(center_freq+actual_bandwidth/2)/1e9:.3f} GHz")
        
        return i_data, q_data
    
    def upload_waveform(self, i_data: np.ndarray, q_data: np.ndarray, 
                       auto_scale: bool = False) -> bool:
        """
        上传波形到SMW仪器
        
        Args:
            i_data: I分量数据
            q_data: Q分量数据
            auto_scale: 是否自动缩放
            
        Returns:
            上传是否成功
        """
        try:
            # 创建临时文件
            with tempfile.NamedTemporaryFile(suffix='.wv', delete=False) as tmp:
                self.pc_wv_file = tmp.name
            
            clock_freq = self.config.get('clock_frequency', 600e6)
            
            # 创建波形文件
            comment = 'FMCW信号 - 未缩放' if not auto_scale else 'FMCW信号 - 自动缩放'
            self.smw.arb_files.create_waveform_file_from_samples(
                i_data, q_data, self.pc_wv_file, 
                clock_freq=clock_freq, 
                auto_scale=auto_scale, 
                comment=comment
            )
            
            self.logger.info(f"波形文件已创建: {self.pc_wv_file}")
            
            # 发送到仪器
            if auto_scale:
                # 上传到Output B
                self.smw.arb_files.send_waveform_file_to_instrument(
                    self.pc_wv_file, self.instr_wv_file_outB
                )
                self.logger.info("波形已上传到Output B (自动缩放)")
            else:
                # 上传到Output A
                self.smw.arb_files.send_waveform_file_to_instrument(
                    self.pc_wv_file, self.instr_wv_file_outA
                )
                self.logger.info("波形已上传到Output A (原始缩放)")
            
            # 清理临时文件
            if os.path.exists(self.pc_wv_file):
                os.remove(self.pc_wv_file)
            
            return True
            
        except Exception as e:
            self.logger.error(f"上传波形失败: {e}")
            return False
    
    def configure_output(self, auto_scale: bool = False) -> bool:
        """
        配置SMW输出参数
        
        Args:
            auto_scale: 是否使用自动缩放的波形
            
        Returns:
            配置是否成功
        """
        try:
            center_freq = self.config.get('center_frequency', 10e9)
            power_level = self.config.get('power_level', 20)
            
            # 选择波形文件
            if auto_scale:
                self.smw.source.bb.arbitrary.waveform.set_select(self.instr_wv_file_outB)
                self.logger.info("已选择Output B波形文件 (自动缩放)")
            else:
                self.smw.source.bb.arbitrary.waveform.set_select(self.instr_wv_file_outA)
                self.logger.info("已选择Output A波形文件 (原始缩放)")
            
            # 设置中心频率
            self.smw.source.frequency.set_frequency(center_freq)
            self.logger.info(f"中心频率已设置为: {center_freq/1e9:.3f} GHz")
            
            # 设置功率电平
            self.smw.source.power.level.immediate.set_amplitude(power_level)
            self.logger.info(f"功率电平已设置为: {power_level} dBm")
            
            # 启用ARB基带
            self.smw.source.bb.arbitrary.set_state(True)
            self.logger.info("ARB基带已启用")
            
            return True
            
        except Exception as e:
            self.logger.error(f"配置输出参数失败: {e}")
            return False
    
    def enable_output(self) -> bool:
        """
        启用RF输出
        
        Returns:
            启用是否成功
        """
        try:
            self.smw.output.state.set_value(True)
            self.logger.info("RF输出已启用")
            return True
            
        except Exception as e:
            self.logger.error(f"启用RF输出失败: {e}")
            return False
    
    def disable_output(self) -> bool:
        """
        禁用RF输出
        
        Returns:
            禁用是否成功
        """
        try:
            #self.smw.output.state.set_value(False)
            self.logger.info("RF输出已禁用")
            return True
            
        except Exception as e:
            self.logger.error(f"禁用RF输出失败: {e}")
            return False
    
    def prepare_fmcw_signal(self, fmcw_config: Dict[str, Any],
                          auto_scale: bool = False) -> bool:
        """
        准备FMCW信号但不立即启用输出
        适用于先准备信号源，再启用输出的工作流程
        
        Args:
            fmcw_config: FMCW信号配置
            auto_scale: 是否使用自动缩放
            
        Returns:
            准备是否成功
        """
        try:
            self.logger.info("=" * 50)
            self.logger.info("准备FMCW信号（输出暂未启用）")
            self.logger.info("=" * 50)
            
            # 生成I/Q数据
            i_data, q_data = self.generate_fmcw_iq_data(fmcw_config)
            
            # 上传波形
            if not self.upload_waveform(i_data, q_data, auto_scale):
                return False
            
            # 配置输出参数（但不启用RF输出）
            if not self.configure_output(auto_scale):
                return False
            
            # 禁用RF输出，等待后续启用
            if not self.disable_output():
                return False
            
            self.logger.info("FMCW信号已准备就绪，等待启用输出...")
            self.logger.info("=" * 50)
            
            return True
            
        except Exception as e:
            self.logger.error(f"准备FMCW信号失败: {e}")
            return False
    
    def generate_fmcw_signal(self, fmcw_config: Dict[str, Any],
                           auto_scale: bool = False,
                           enable_immediately: bool = True) -> bool:
        """
        完整的FMCW信号生成流程
        
        Args:
            fmcw_config: FMCW信号配置
            auto_scale: 是否使用自动缩放
            enable_immediately: 是否立即启用输出
            
        Returns:
            生成是否成功
        """
        try:
            self.logger.info("=" * 50)
            if enable_immediately:
                self.logger.info("开始生成FMCW信号")
            else:
                self.logger.info("准备FMCW信号（不立即启用输出）")
            self.logger.info("=" * 50)
            
            # 生成I/Q数据
            i_data, q_data = self.generate_fmcw_iq_data(fmcw_config)
            
            # 上传波形
            if not self.upload_waveform(i_data, q_data, auto_scale):
                return False
            
            # 配置输出
            if not self.configure_output(auto_scale):
                return False
            
            # 根据参数决定是否立即启用输出
            if enable_immediately:
                if not self.enable_output():
                    return False
                self.logger.info("=" * 50)
                self.logger.info("FMCW信号生成完成")
            else:
                if not self.disable_output():
                    return False
                self.logger.info("FMCW信号准备完成，等待启用输出")
            self.logger.info("=" * 50)
            
            return True
            
        except Exception as e:
            self.logger.error(f"生成FMCW信号失败: {e}")
            return False
    
    def enable_output_with_delay(self, delay_seconds: float = 0.0) -> bool:
        """
        延迟启用RF输出
        
        Args:
            delay_seconds: 延迟时间（秒）
            
        Returns:
            启用是否成功
        """
        try:
            if delay_seconds > 0:
                self.logger.info(f"等待 {delay_seconds} 秒后启用输出...")
                import time
                time.sleep(delay_seconds)
            
            return self.enable_output()
            
        except Exception as e:
            self.logger.error(f"延迟启用输出失败: {e}")
            return False
    
    def get_instrument_info(self) -> Dict[str, Any]:
        """
        获取仪器信息
        
        Returns:
            仪器信息字典
        """
        if self.smw is None:
            return {}
        
        try:
            return {
                'idn': self.smw.utilities.idn_string,
                'driver_version': self.smw.utilities.driver_version,
                'connected': True
            }
        except Exception:
            return {'connected': False}
    
    def disconnect(self) -> None:
        """断开与SMW的连接"""
        try:
            if self.smw:
                # 禁用输出
                self.disable_output()
                
                # 关闭连接
                self.smw.close()
                self.smw = None
                self.logger.info("SMW信号发生器连接已关闭")
                
        except Exception as e:
            self.logger.error(f"断开SMW连接时出错: {e}")
    
    def __del__(self):
        """析构函数，确保连接被正确关闭"""
        self.disconnect()